#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
市场特征和基金数据获取功能模块

包含以下功能：
1. 市场指数数据获取（上证、深证、沪深300）
2. 基金净值数据获取
3. 市场特征添加功能
4. A股全市场交易量数据获取（市场情绪指标）
5. 市场特征缓存机制

作者：AI Assistant
创建时间：2024年
"""

import requests
import pandas as pd
import numpy as np
from datetime import datetime, timedelta
from tqdm import tqdm
import time
import pickle
import os
import json


def get_market_volume_data(date_str):
    """获取A股全市场交易量数据，用于衡量市场情绪
    
    Args:
        date_str: 日期字符串，格式为 YYYY-MM-DD
    
    Returns:
        dict: 包含沪深两市总成交量、成交额等数据
    """
    import requests
    import json
    from datetime import datetime
    
    # 转换日期格式
    try:
        date_obj = datetime.strptime(date_str, '%Y-%m-%d')
        date_formatted = date_obj.strftime('%Y%m%d')
    except:
        return {
            'total_volume': None,
            'total_amount': None,
            'sh_volume': None,
            'sh_amount': None,
            'sz_volume': None,
            'sz_amount': None,
            'market_sentiment_score': None,
            'is_trading_day': 0
        }
    
    headers = {
        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36',
        'Referer': 'https://quote.eastmoney.com/'
    }
    
    # 上证指数 (000001.SH)、深证成指 (399001.SZ)
    index_codes = {
        'sh': '1.000001',      # 上证指数
        'sz': '0.399001',      # 深证成指  
    }
    
    result = {
        'total_volume': None,
        'total_amount': None,
        'sh_volume': None,
        'sh_amount': None,
        'sz_volume': None,
        'sz_amount': None,
        'market_sentiment_score': None,
        'is_trading_day': 0
    }
    
    try:
        # 获取沪深两市成交量成交额数据
        for key, code in index_codes.items():
            # 使用历史数据接口获取指定日期的数据，包含成交量和成交额
            hist_url = f"https://push2his.eastmoney.com/api/qt/stock/kline/get?secid={code}&fields1=f1,f2,f3,f4,f5,f6&fields2=f51,f52,f53,f54,f55,f56,f57,f58,f59,f60,f61&klt=101&fqt=1&end={date_formatted}&lmt=1"
            
            try:
                response = requests.get(hist_url, headers=headers, timeout=10)
                if response.status_code == 200:
                    data = response.json()
                    if data.get('data') and data['data'].get('klines'):
                        kline = data['data']['klines'][0].split(',')
                        if len(kline) >= 6:
                            # kline数据格式：日期,开盘,收盘,最高,最低,成交量,成交额,振幅,涨跌幅,涨跌额,换手率
                            volume = float(kline[5]) if kline[5] else 0  # 成交量（手）
                            amount = float(kline[6]) if kline[6] else 0  # 成交额（元）
                            
                            result[f'{key}_volume'] = volume
                            result[f'{key}_amount'] = amount
                            result['is_trading_day'] = 1
                            
            except Exception as e:
                print(f"获取{key}交易量数据失败: {e}")
                continue
        
        # 计算总成交量和总成交额
        if result['sh_volume'] is not None and result['sz_volume'] is not None:
            result['total_volume'] = result['sh_volume'] + result['sz_volume']
        elif result['sh_volume'] is not None:
            result['total_volume'] = result['sh_volume']
        elif result['sz_volume'] is not None:
            result['total_volume'] = result['sz_volume']
            
        if result['sh_amount'] is not None and result['sz_amount'] is not None:
            result['total_amount'] = result['sh_amount'] + result['sz_amount']
        elif result['sh_amount'] is not None:
            result['total_amount'] = result['sh_amount']
        elif result['sz_amount'] is not None:
            result['total_amount'] = result['sz_amount']
        
        # 计算市场情绪评分（基于成交量）
        # 这是一个简单的情绪指标，可以根据需要调整算法
        if result['total_volume'] is not None:
            # 将成交量转换为亿手为单位，便于计算
            volume_yi = result['total_volume'] / 100000000
            
            # 简单的情绪评分：成交量越大，市场情绪越活跃
            # 可以根据历史数据调整这些阈值
            if volume_yi >= 200:
                result['market_sentiment_score'] = 'very_high'  # 极度活跃
            elif volume_yi >= 150:
                result['market_sentiment_score'] = 'high'       # 高度活跃
            elif volume_yi >= 100:
                result['market_sentiment_score'] = 'medium'     # 中等活跃
            elif volume_yi >= 50:
                result['market_sentiment_score'] = 'low'        # 低度活跃
            else:
                result['market_sentiment_score'] = 'very_low'   # 极度低迷
                
    except Exception as e:
        print(f"获取市场交易量数据失败: {e}")
    
    return result


def get_market_index_data(date_str):
    """获取指定日期的市场指数数据
    
    Args:
        date_str: 日期字符串，格式为 YYYY-MM-DD
    
    Returns:
        dict: 包含上证指数、深证成指、沪深300的数据
    """
    import requests
    import json
    from datetime import datetime
    
    # 转换日期格式
    try:
        date_obj = datetime.strptime(date_str, '%Y-%m-%d')
        date_formatted = date_obj.strftime('%Y%m%d')
    except:
        return {
            'sh_index': None, 'sh_change': None,
            'sz_index': None, 'sz_change': None, 
            'hs300_index': None, 'hs300_change': None,
            'is_trading_day': 0
        }
    
    headers = {
        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36',
        'Referer': 'https://quote.eastmoney.com/'
    }
    
    # 上证指数 (000001.SH)、深证成指 (399001.SZ)、沪深300 (000300.SH)
    index_codes = {
        'sh': '1.000001',      # 上证指数
        'sz': '0.399001',      # 深证成指  
        'hs300': '1.000300'    # 沪深300
    }
    
    result = {
        'sh_index': None, 'sh_change': None,
        'sz_index': None, 'sz_change': None,
        'hs300_index': None, 'hs300_change': None,
        'is_trading_day': 0
    }
    
    try:
        # 获取指数实时数据
        for key, code in index_codes.items():
            url = f"https://push2.eastmoney.com/api/qt/stock/get?secid={code}&fields=f43,f44,f45,f46,f47,f48,f49,f50,f51,f52,f53,f54,f55,f56,f57,f58,f170"
            
            try:
                response = requests.get(url, headers=headers, timeout=10)
                if response.status_code == 200:
                    data = response.json()
                    if data.get('data'):
                        stock_data = data['data']
                        current_price = stock_data.get('f43')  # 当前价格
                        change_pct = stock_data.get('f170')    # 涨跌幅
                        
                        if current_price:
                            result[f'{key}_index'] = float(current_price) / 100 if current_price else None
                            result[f'{key}_change'] = float(change_pct) / 10000 if change_pct else None  # f170字段是基点形式
                            result['is_trading_day'] = 1
                            
            except Exception as e:
                print(f"获取{key}指数数据失败: {e}")
                continue
        
        # 使用历史数据接口获取指定日期的准确数据
        for key, code in index_codes.items():
            if not result.get(f'{key}_index'):  # 如果实时数据没获取到
                hist_url = f"https://push2his.eastmoney.com/api/qt/stock/kline/get?secid={code}&fields1=f1,f2,f3,f4,f5,f6&fields2=f51,f52,f53,f54,f55,f56,f57,f58,f59,f60,f61&klt=101&fqt=1&end={date_formatted}&lmt=1"
                
                try:
                    response = requests.get(hist_url, headers=headers, timeout=10)
                    if response.status_code == 200:
                        data = response.json()
                        if data.get('data') and data['data'].get('klines'):
                            kline = data['data']['klines'][0].split(',')
                            if len(kline) >= 9:
                                result[f'{key}_index'] = float(kline[2])  # 收盘价
                                result[f'{key}_change'] = float(kline[8]) / 100  # 涨跌幅（东财返回的是百分比，需要除以100转为小数）
                                result['is_trading_day'] = 1
                except Exception as e:
                    print(f"获取{key}历史数据失败: {e}")
                    continue
                
    except Exception as e:
        print(f"获取市场指数数据失败: {e}")
    
    return result


def get_fund_nav_data(fund_code, date_str):
    """获取基金净值数据
    
    Args:
        fund_code: 基金代码
        date_str: 日期字符串，格式为 YYYY-MM-DD
    
    Returns:
        dict: 包含基金净值和涨跌幅数据
    """
    import requests
    from datetime import datetime, timedelta
    
    # 确保基金代码为6位字符串
    fund_code = str(fund_code).zfill(6)
    
    headers = {
        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36',
        'Referer': 'https://fundf10.eastmoney.com/'
    }
    
    result = {
        'fund_nav': None,
        'fund_nav_change': None,
        'fund_accumulated_nav': None
    }
    
    try:
        # 获取基金净值数据
        url = f"https://api.fund.eastmoney.com/f10/lsjz?fundCode={fund_code}&pageIndex=1&pageSize=10"
        
        response = requests.get(url, headers=headers, timeout=10)
        if response.status_code == 200:
            data = response.json()
            if data.get('Data') and data['Data'].get('LSJZList'):
                nav_list = data['Data']['LSJZList']
                
                # 查找指定日期的净值数据
                target_date = datetime.strptime(date_str, '%Y-%m-%d').strftime('%Y-%m-%d')
                
                for nav_item in nav_list:
                    if nav_item.get('FSRQ') == target_date:
                        result['fund_nav'] = float(nav_item['DWJZ']) if nav_item.get('DWJZ') else None
                        result['fund_nav_change'] = float(nav_item['JZZZL']) / 100 if nav_item.get('JZZZL') else None
                        result['fund_accumulated_nav'] = float(nav_item['LJJZ']) if nav_item.get('LJJZ') else None
                        break
                
                # 如果没找到当天数据，使用最近的数据
                if not result['fund_nav'] and nav_list:
                    latest_nav = nav_list[0]
                    result['fund_nav'] = float(latest_nav['DWJZ']) if latest_nav.get('DWJZ') else None
                    result['fund_nav_change'] = float(latest_nav['JZZZL']) / 100 if latest_nav.get('JZZZL') else None
                    result['fund_accumulated_nav'] = float(latest_nav['LJJZ']) if latest_nav.get('LJJZ') else None
                    
    except Exception as e:
        print(f"获取基金{fund_code}净值数据失败: {e}")
    
    return result


def load_market_features_cache(cache_path="market_features_cache.pkl"):
    """加载市场特征缓存
    
    Args:
        cache_path: 缓存文件路径
    
    Returns:
        dict: 缓存的市场数据，格式为 {date_str: {market_data, volume_data}}
    """
    if os.path.exists(cache_path):
        try:
            with open(cache_path, 'rb') as f:
                cache_data = pickle.load(f)
            print(f"✅ 加载市场特征缓存: {cache_path}")
            print(f"📊 缓存包含 {len(cache_data.get('market_cache', {}))} 个交易日的市场数据")
            print(f"📊 缓存包含 {len(cache_data.get('volume_cache', {}))} 个交易日的交易量数据")
            return cache_data
        except Exception as e:
            print(f"⚠️ 缓存加载失败: {e}")
            return {'market_cache': {}, 'volume_cache': {}, 'timestamp': None}
    else:
        print(f"📁 缓存文件不存在，将创建新的缓存: {cache_path}")
        return {'market_cache': {}, 'volume_cache': {}, 'timestamp': None}

def save_market_features_cache(market_cache, volume_cache, cache_path="market_features_cache.pkl"):
    """保存市场特征缓存
    
    Args:
        market_cache: 市场指数数据缓存
        volume_cache: 交易量数据缓存
        cache_path: 缓存文件路径
    """
    try:
        cache_data = {
            'market_cache': market_cache,
            'volume_cache': volume_cache,
            'timestamp': datetime.now().isoformat(),
            'total_dates': len(market_cache)
        }
        
        with open(cache_path, 'wb') as f:
            pickle.dump(cache_data, f)
        
        print(f"💾 市场特征缓存已保存: {cache_path}")
        print(f"📊 保存了 {len(market_cache)} 个交易日的市场数据")
        print(f"📊 保存了 {len(volume_cache)} 个交易日的交易量数据")
        
    except Exception as e:
        print(f"❌ 缓存保存失败: {e}")

def get_missing_dates(required_dates, cached_dates):
    """获取需要新获取的日期列表
    
    Args:
        required_dates: 需要的日期列表
        cached_dates: 已缓存的日期列表
    
    Returns:
        list: 需要新获取的日期列表
    """
    required_set = set(required_dates)
    cached_set = set(cached_dates)
    missing_dates = list(required_set - cached_set)
    
    if missing_dates:
        print(f"📅 需要新获取 {len(missing_dates)} 个日期的市场数据")
        print(f"📅 已缓存 {len(cached_set & required_set)} 个日期的数据")
    else:
        print(f"✅ 所有需要的日期数据都已缓存")
    
    return missing_dates

def add_market_features(df, cache_path="market_features_cache.pkl", force_refresh=False):
    """为数据添加市场特征、基金净值特征和市场交易量特征（带缓存机制）
    
    Args:
        df: 输入数据框
        cache_path: 市场特征缓存文件路径
        force_refresh: 是否强制刷新缓存（忽略现有缓存）
    
    Returns:
        df: 添加了市场特征的数据框
    """
    print("🔄 添加市场指数、基金净值和市场交易量特征...")
    
    # 初始化新列
    market_cols = ['sh_index', 'sh_change', 'sz_index', 'sz_change', 
                   'hs300_index', 'hs300_change', 'is_trading_day']
    fund_cols = ['fund_nav', 'fund_nav_change', 'fund_accumulated_nav']
    volume_cols = ['total_volume', 'total_amount', 'sh_volume', 'sh_amount',
                   'sz_volume', 'sz_amount', 'market_sentiment_score']
    
    for col in market_cols + fund_cols + volume_cols:
        df[col] = None
    
    # 获取唯一的日期列表
    unique_dates = df['transaction_date'].dt.strftime('%Y-%m-%d').unique().tolist()
    print(f"📊 需要处理 {len(unique_dates)} 个唯一交易日")
    
    # 加载现有缓存
    if force_refresh:
        print("🔄 强制刷新模式，忽略现有缓存")
        cache_data = {'market_cache': {}, 'volume_cache': {}, 'timestamp': None}
    else:
        cache_data = load_market_features_cache(cache_path)
    
    market_cache = cache_data.get('market_cache', {})
    volume_cache = cache_data.get('volume_cache', {})
    
    # 确定需要新获取的日期
    missing_dates = get_missing_dates(unique_dates, list(market_cache.keys()))
    
    # 获取缺失的市场数据
    if missing_dates:
        print(f"🌐 开始获取 {len(missing_dates)} 个日期的新市场数据...")
        
        for date_str in tqdm(missing_dates, desc="获取新市场数据"):
            # 获取市场指数数据
            market_data = get_market_index_data(date_str)
            market_cache[date_str] = market_data
            
            # 获取交易量数据
            volume_data = get_market_volume_data(date_str)
            volume_cache[date_str] = volume_data
            
            # 添加延迟避免请求过快
            time.sleep(0.1)
        
        # 保存更新后的缓存
        save_market_features_cache(market_cache, volume_cache, cache_path)
    
    # 为每行数据添加市场特征
    print("🔗 开始拼接市场特征到数据...")
    
    for idx, row in tqdm(df.iterrows(), total=len(df), desc="添加市场特征"):
        date_str = row['transaction_date'].strftime('%Y-%m-%d')
        
        # 添加市场指数数据
        if date_str in market_cache:
            market_data = market_cache[date_str]
            for col in market_cols:
                df.at[idx, col] = market_data.get(col)
        else:
            print(f"⚠️ 日期 {date_str} 的市场数据不在缓存中")
        
        # 添加交易量数据
        if date_str in volume_cache:
            volume_data = volume_cache[date_str]
            for col in volume_cols:
                df.at[idx, col] = volume_data.get(col)
        else:
            print(f"⚠️ 日期 {date_str} 的交易量数据不在缓存中")
        
        # 获取基金净值数据（基金数据暂不缓存，因为每个基金代码都不同）
        fund_code = row['fund_code']
        fund_data = get_fund_nav_data(fund_code, date_str)
        for col in fund_cols:
            df.at[idx, col] = fund_data.get(col)
        
        # 每处理100行添加一个延迟
        if idx % 100 == 0:
            time.sleep(0.05)
    
    # 统计结果
    total_features = len(market_cols + fund_cols + volume_cols)
    print(f"✅ 特征添加完成！")
    print(f"📊 数据从 {len(df.columns) - total_features} 列扩展到 {len(df.columns)} 列")
    print(f"📊 新增市场特征: {len(market_cols)} 个")
    print(f"📊 新增基金特征: {len(fund_cols)} 个") 
    print(f"📊 新增交易量特征: {len(volume_cols)} 个")
    
    return df

def clear_market_features_cache(cache_path="market_features_cache.pkl"):
    """清空市场特征缓存
    
    Args:
        cache_path: 缓存文件路径
    """
    try:
        if os.path.exists(cache_path):
            os.remove(cache_path)
            print(f"🗑️ 市场特征缓存已清空: {cache_path}")
        else:
            print(f"📁 缓存文件不存在: {cache_path}")
    except Exception as e:
        print(f"❌ 缓存清空失败: {e}")

def get_cache_info(cache_path="market_features_cache.pkl"):
    """获取缓存信息
    
    Args:
        cache_path: 缓存文件路径
    
    Returns:
        dict: 缓存信息
    """
    if not os.path.exists(cache_path):
        return {"exists": False, "message": "缓存文件不存在"}
    
    try:
        with open(cache_path, 'rb') as f:
            cache_data = pickle.load(f)
        
        market_cache = cache_data.get('market_cache', {})
        volume_cache = cache_data.get('volume_cache', {})
        timestamp = cache_data.get('timestamp', 'Unknown')
        
        # 获取日期范围
        if market_cache:
            dates = list(market_cache.keys())
            dates.sort()
            date_range = f"{dates[0]} 到 {dates[-1]}"
        else:
            date_range = "无数据"
        
        return {
            "exists": True,
            "timestamp": timestamp,
            "market_dates_count": len(market_cache),
            "volume_dates_count": len(volume_cache),
            "date_range": date_range,
            "file_size": f"{os.path.getsize(cache_path) / 1024:.2f} KB"
        }
        
    except Exception as e:
        return {"exists": True, "error": str(e)}
